# -*- coding: UTF-8 -*-

from pymongo import MongoClient
from WindPy import *

import time
import datetime
import math



class DailyClose:
    def __init__(self, startDateStr, endDateStr = time.strftime('%Y-%m-%d'), fillType='F', DBIP='localhost', DBPort=27017, DBName = 'stocks', collectionPreName = 'daily_close', txtFileName = 'close.txt'):
        self.startDateStr = startDateStr
        self.endDateStr = endDateStr
        self.fillType = fillType
        self.DBIP = DBIP
        self.DBPort = DBPort
        self.DBName = DBName
        self.collectionPreName = collectionPreName
        self.txtFileName = txtFileName

    def writeFromWindToMongoDB(self):
        # format of tje default 'startDateStr' is '2012-01-09'
        # format of tje default 'endDateStr' is '2017-05-08'
        # fillType = 'F': 前复权 ，fillType = 'B': 后复权

        startDateStr = self.startDateStr
        endDateStr = self.endDateStr
        fillType = self.fillType
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        collectionPreName = self.collectionPreName

        w.start()

        resStockID = w.wset('sectorconstituent', 'date=' + endDateStr + ';sectorid=a001010100000000;field=wind_code')

        stockIDdata = resStockID.Data[0]

        totalStockIDnum = len(stockIDdata)

        stockIdList = []

        for i in range(0, totalStockIDnum):
            stockIdStr = str(stockIDdata[i])

            stockIdList.append(stockIdStr)

        stockIndexIDList = ['000001.SH', '399001.SZ', '399006.SZ', '000300.SH', '000016.SH', '000905.SH', '399005.SZ',
                            '881001.WI']

        stockIdList = stockIdList + stockIndexIDList

        totalStockIDnum = len(stockIdList)

        client = MongoClient(DBIP, DBPort)

        db = client[DBName]

        collectionName = collectionPreName + '_' + startDateStr + '-' + endDateStr + '_' + fillType

        collection_dailyclose = db[collectionName]

        for stockNdx in range(0, totalStockIDnum):
            oneStockID = stockIdList[stockNdx]
            resDailyClose = w.wsd(oneStockID, 'close', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)
            closeDataNum = resDailyClose.Data[0]
            closeDataTime = resDailyClose.Times
            totalTransDaysNum = len(closeDataNum)
            aData = []
            aDateDict = {}
            aDateDict['StockID'] = oneStockID
            for tDayNum in range(0, totalTransDaysNum):
                aDateDict[closeDataTime[tDayNum].strftime('%Y%m%d')] = closeDataNum[tDayNum]

            aData.append(aDateDict)
            collection_dailyclose.insert(aData)

            print('Daily close data of ' + oneStockID + ' wrote to MongoDB...')
        print('Finished writing close data of all stocks to MongoDB ')


    def writeFromMongoDBtoTxt(self):

        # format of tje default 'startDateStr' is '2012-01-09'
        # format of tje default 'endDateStr' is '2017-05-08'
        # fillType = 'F': 前复权 ，fillType = 'B': 后复权

        startDateStr = self.startDateStr
        endDateStr = self.endDateStr
        fillType = self.fillType
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        collectionPreName = self.collectionPreName
        txtFileName = self.txtFileName

        client = MongoClient(DBIP, DBPort)

        db = client[DBName]

        collectionName = collectionPreName + '_' + startDateStr + '-' + endDateStr + '_' + fillType

        collection_dailyclose = db[collectionName]

        dateList = []
        totalStockCloseList = []

        dateListRes = collection_dailyclose.find({'StockID': '000001.SH'})
        dateListDict = dateListRes[0]

        for (k, v) in dateListDict.items():
            dateStr = ''
            try:
                dateStr = datetime.datetime.strptime(k, '%Y-%m-%d').strftime('%Y%m%d')
            except:
                dateStr = k
            dateList.append(dateStr)
        dateList = dateList[2:]
        dateList.insert(0, 'close')
        totalStockCloseList.append(dateList)

        for stockIdRes in collection_dailyclose.find({}, {'StockID': 1}):
            stockIdStr = stockIdRes['StockID']
            for stockDailyCloseResDict in collection_dailyclose.find({'StockID': stockIdStr}):
                oneStockCloseList = []
                for i in stockDailyCloseResDict:
                    oneStockCloseList.append(stockDailyCloseResDict[i])
                oneStockCloseList = oneStockCloseList[1:]
                totalStockCloseList.append(oneStockCloseList)
        print('Finished retrieving close data from MongoDB...')
        f = open(txtFileName, 'w')

        txtFileRowNum = len(oneStockCloseList)
        txtFileColNum = len(totalStockCloseList)

        f.write(str(txtFileRowNum - 1) + '\t' + str(txtFileColNum - 1) + '\t')

        for firstRowLen in range(txtFileColNum - 1):
            f.write('\t')
        f.write('\n')
        print('Finished writing text file header...')

        for rowIter in range(txtFileRowNum):
            for colIter in range(txtFileColNum):
                try:
                    cellData = totalStockCloseList[colIter][rowIter]
                    if math.isnan(float(cellData)):
                        cellData = 0
                    f.write(str(cellData) + '\t')
                except ValueError:
                    cellData = str(totalStockCloseList[colIter][rowIter])
                    f.write(cellData + '\t')

            f.write('\n')
            print('Wrote row number # ' + str(rowIter) + ' to the text file....')

        f.close()
        print('Finished writing the whole text file!')

obj = DailyClose('2013-01-01')

obj.writeFromWindToMongoDB()
obj.writeFromMongoDBtoTxt()